//      build_vom.cc
//      
//      Copyright 2012 tqlong <tqlong@espada>
//      
//      This program is free software; you can redistribute it and/or modify
//      it under the terms of the GNU General Public License as published by
//      the Free Software Foundation; either version 2 of the License, or
//      (at your option) any later version.
//      
//      This program is distributed in the hope that it will be useful,
//      but WITHOUT ANY WARRANTY; without even the implied warranty of
//      MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//      GNU General Public License for more details.
//      
//      You should have received a copy of the GNU General Public License
//      along with this program; if not, write to the Free Software
//      Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
//      MA 02110-1301, USA.

#include <iostream>
#include <boost/program_options.hpp>
#include <boost/date_time/posix_time/posix_time.hpp> 

#include <boost/foreach.hpp> 
#include <map>
#include "stl-helper.h"

#include "dna_processing.h"
#include "vom_builder.h"

namespace po = boost::program_options;
po::options_description desc("Allowed options");
po::variables_map vm;

void process_options(int argc, char** argv);
void usage(int exitcode = 0);

void read_fasta();
void suffix_array_test();
void build_fom();
void likelihood_compute();
void grow_vom();
void simulate();
void rand_fom();
void accuracy();
void depth();
void split();
void split_train_test();
void fix_cds();
void gc_content();

int main(int argc, char** argv)
{
  //MSG("Variable Order Model ... ");

  process_options(argc, argv);
  srand(time(NULL));

  boost::posix_time::ptime time_start(boost::posix_time::microsec_clock::local_time());

  if (vm.count("help"))
    usage();
  else if (vm.count("read_fasta"))
    read_fasta();
  else if (vm.count("suffix_array"))
    suffix_array_test();
  else if (vm.count("fom"))
    build_fom();
  else if (vm.count("likelihood"))
    likelihood_compute();
  else if (vm.count("grow"))
    grow_vom();
  else if (vm.count("simulate"))
    simulate();
  else if (vm.count("rand_fom"))
    rand_fom();
  else if (vm.count("accuracy"))
    accuracy();
  else if (vm.count("depth"))
    depth();
  else if (vm.count("split"))
    split();
  else if (vm.count("train_test"))
    split_train_test();
  else if (vm.count("fix_cds"))
    fix_cds();
  else if (vm.count("gc"))
    gc_content();
  else
    usage();

  boost::posix_time::ptime time_end(boost::posix_time::microsec_clock::local_time());
  boost::posix_time::time_duration duration(time_end - time_start);
  MSG("Duration = " << duration);
}

void process_options(int argc, char** argv)
{
  desc.add_options()
    ("help", "produce this usage")
    ("input", po::value<std::string>(), "input file")
    ("inputs", po::value<std::vector<std::string> >(), "input files")
    ("output", po::value<std::string>(), "output file")
    ("model", po::value<std::string>(), "model file")
    ("seed", po::value<int>(), "random seed (REQUIRED)")

    // PARAMETERS
    ("order", po::value<int>(), "order of the model")
    ("pattern", po::value<std::string>(), "pattern to search for")
    ("eta", po::value<double>()->default_value(0.1), "KL threshold (default 0.1)")
    ("n", po::value<int>(), "number of sequences")
    ("length", po::value<int>(), "sequence length")
    ("maxlen", po::value<int>(), "maximum sequence length")
    ("period", po::value<int>()->default_value(1), "the periodicity of the model (1,2,3)")
    ("ll", po::value<std::vector<std::string> >(), "sequence files to compute log-likelihood")
    ("llout", po::value<std::vector<std::string> >(), "likelihood output files")
    ("depthout", po::value<std::string>(), "stats output file")
    ("prunefirst", "prune before grow")

    // Accuracy calculation
    ("negneg", po::value<std::string>(), "negative set, negative model output")
    ("negpos", po::value<std::string>(), "negative set, positive model output")
    ("posneg", po::value<std::string>(), "positive set, negative model output")
    ("pospos", po::value<std::string>(), "positive set, positive model output")

    // Train test split
    ("train", po::value<std::string>(), "train set fasta file")
    ("test", po::value<std::string>(), "test set fasta file")
    ("ratio", po::value<double>()->default_value(0.5), "ratio of train set / all")

    //JOB TYPE
    ("read_fasta", "read fasta file\n--input=<FASTA FILE>")
    ("suffix_array", "suffix array test\n--input=<FASTA FILE>")
    ("fom", "FOM build\n--input=<FASTA FILE> --output=<MODEL FILE> --order=<ORDER>")
    ("grow", po::value<std::string>(), "grow vom tree\n--grow=[grow,growprune,prunegrow,prune] --input=<FASTA FILE> --model=<MODEL FILE> --output=<GROWED MODEL FILE> [--prune] [--eta]")
    ("prune", "prune after grow")
    ("likelihood", "likelihood computation\n--input=<FASTA FILE> --model=<MODEL FILE> --output=")
    ("simulate", "generate simulated sequence\n--model=<MODEL FILE> --output=<OUTPUT FASTA FILE> --n=<NUMBER OF SEQUENCES> --length=<SEQUENCE LENGTH>")
    ("rand_fom", "Random FOM build\n--output=<MODEL FILE> --order=<ORDER>")    
    ("accuracy", "--negneg --negpos --pospos --posneg --output")
    ("depth", "calculate depth of a model\n--model=<MODEL FILE>")
    ("split", "split long sequence to fixed length\n --input --output --length")
    ("train_test", "split set of sequences to train set and test seth\n --input --train --test --ratio")
    ("fix_cds", "fix cds sequences\n --input [--output]")
    ("gc", "calculate GC-content\n [--input] [--inputs]")
    ;
  po::store(po::parse_command_line(argc, argv, desc), vm);
  po::notify(vm);
}

void usage(int exitcode)
{
  std::cout << desc << std::endl;
  exit(exitcode);
}

void read_fasta()
{
  std::string inputFile = vm["input"].as<std::string>();
  Sequences seqs;

  MSG("Reading fasta file  ... " << inputFile);
  seqs.fasta_read(inputFile);

  if (!vm.count("output")) {
    if (vm.count("verbose")) seqs.print(std::cerr);
    MSG("n_seqs = " << sz(seqs));
  }
  else {
    std::string outputFile = vm["output"].as<std::string>();
    FileAccessWrite(output, outputFile);
    seqs.print(output, true);
    output.close();
  }
}

void suffix_array_test() {
  vom::sastring str = VM(vm, "input", std::string);
  std::string p = VM(vm, "pattern", std::string);
  std::string output = VM(vm, "output", std::string);
  int order = VM(vm, "order", int);
  int period = VM(vm, "period", int);

  for (uint i = 0; i < str.length(); i++)
    MSG(i << " " << str.SA[i] << " " << str.substr(str.SA[i]));
  MSG("result=" << str.result);
  
  int left, size;
  str.sasearch(p, left, size);
  MSG("left=" << left << " size=" << size);

  vom::VOM m("ACGT", period);
  m.build_fom(str, order);
  MSG("M = ");
  //m.print_node(std::cerr, m.root_, 0);
	
  FileAccessWrite(out, output);
  m.print_node(out);
  out.close();
	
  FileAccessRead(in, output);
  vom::VOM m_copy("ACGT", period);
  m_copy.load_node(in);
  MSG("M_COPY = ");
  //m_copy.print_node(std::cerr, m_copy.root_, 0);
}

void build_fom() {
  std::string input = VM(vm, "input", std::string);
  int order = VM(vm, "order", int);
  int period = VM(vm, "period", int);
  MSG("++ Build FOM input=" << input << " order=" << order);
	
  Sequences seqs;
  seqs.fasta_read(input);
	
  std::vector<vom::sastring> train_set;
  BOOST_FOREACH(const std::string& s, seqs)
    train_set.pb(s);
		
  vom::VOM m("ACGT", period);
  m.build_fom(train_set, order);
	
  if (vm.count("output")) {
    std::string output = VM(vm, "output", std::string);
    FileAccessWrite(out, output);
    m.print_node(out);
  }

  if (vm.count("ll") && vm.count("llout")) {
    std::vector<std::string> ll = VM(vm, "ll", std::vector<std::string>);
    std::vector<std::string> llout = VM(vm, "llout", std::vector<std::string>);
    if (sz(ll) != sz(llout)) 
      FATAL_ERROR("Check --ll and --llout, should have same number of file names");
    for (int i = 0; i < sz(ll); i++) {
      MSG("Likelihood on " << ll.at(i) << " to " << llout.at(i));
      Sequences seqs;
      seqs.fasta_read(ll.at(i));

      FileAccessWrite(out, llout.at(i));
      int j = 0;
      BOOST_FOREACH(const std::string& s, seqs) {
	double ll = m.log_prob(s);
	out << ">" << seqs.annotation_[j] << ENDL << ll << ENDL;
	j++;
      }
    }
  }
  
  if (vm.count("depthout")) {
    m.depth();
    FileAccessWrite(out, VM(vm, "depthout", std::string));
    out << m.count_ << "\t" << m.depth_  << "\t" << m.count_ << "\t" << m.depth_ << "\t" << m.count_ << "\t" << m.depth_ << "\t" << m.count_ << "\t" << m.depth_ << ENDL;    
  }
}

void likelihood_compute()
{
  std::string input = VM(vm, "input", std::string);
  std::string output = VM(vm, "output", std::string);
  std::string model = VM(vm, "model", std::string);
  int period = VM(vm, "period", int);
  MSG("++ Compute likelihood input=" << input << " output=" << output 
      << " model=" << model);
	
  FileAccessRead(in, model);
  vom::VOM m("ACGT", period);
  m.load_node(in);
	
  Sequences seqs;
  seqs.fasta_read(input);
  MSG("n_seq="<<sz(seqs));
	
  FileAccessWrite(out, output);
  int i = 0;
  BOOST_FOREACH(const std::string& s, seqs) {
    double ll = m.log_prob(s);
    out << ">" << seqs.annotation_[i] << ENDL << ll << ENDL;
    i++;
  }
}

void grow_vom()
{
  std::string input = VM(vm, "input", std::string);
  std::string model = VM(vm, "model", std::string);
  std::string grow = VM(vm, "grow", std::string);
  if (grow != "grow" && grow != "growprune" && grow != "prune" && grow != "prunegrow") {
    MSG("ERROR: check --grow");
    usage();
  }
  
  double eta = VM(vm, "eta", double);
  int period = VM(vm, "period", int);
  MSG("++ Grow VOM input=" << input << " mode=" << grow
      << " model=" << model << " eta=" << eta);
	
  FileAccessRead(in, model);
  vom::VOM m("ACGT", period);
  m.load_node(in);
  MSG("n_node=" << m.count_);
	
  Sequences seqs;
  std::vector<vom::sastring> train_set;

  seqs.fasta_read(input);
  BOOST_FOREACH(const std::string& s, seqs)
    train_set.pb(s);

  if (grow == "grow") {
    m.grow(train_set, eta);
  }
  else if (grow == "growprune") {
    m.grow(train_set, eta);
    double eta_prune = eta / 10;
    if (vm.count("eta_prune")) eta_prune = VM(vm, "eta_prune", double);
    m.prune(train_set, eta_prune);
  }
  else if (grow == "prune") {
    m.prune(train_set, eta);
  }
  else {
    m.prune(train_set, eta);
    m.grow(train_set, eta);
  }
  MSG("after n_node=" << m.count_);

  if (vm.count("output")) {
    std::string output = VM(vm, "output", std::string);
    FileAccessWrite(out, output);
    m.print_node(out);
  }

  if (vm.count("ll") && vm.count("llout")) {
    std::vector<std::string> ll = VM(vm, "ll", std::vector<std::string>);
    std::vector<std::string> llout = VM(vm, "llout", std::vector<std::string>);
    if (sz(ll) != sz(llout)) 
      FATAL_ERROR("Check --ll and --llout, should have same number of file names");
    for (int i = 0; i < sz(ll); i++) {
      MSG("Likelihood on " << ll.at(i) << " to " << llout.at(i));
      Sequences seqs;
      seqs.fasta_read(ll.at(i));

      FileAccessWrite(out, llout.at(i));
      int j = 0;
      BOOST_FOREACH(const std::string& s, seqs) {
	double ll = m.log_prob(s);
	out << ">" << seqs.annotation_[j] << ENDL << ll << ENDL;
	j++;
      }
    }
  }
  
  if (vm.count("depthout")) {
    m.depth();
    FileAccessWrite(out, VM(vm, "depthout", std::string));
    out << m.count_ << "\t" << m.depth_;
    for (int pr = 0; pr < period; pr++)
      out << "\t" << m.count_p_[pr] << "\t" << m.depth_p_[pr];
    out << ENDL;
  }
}

void simulate()
{
  std::string output = VM(vm, "output", std::string);
  std::string model = VM(vm, "model", std::string);
  int n = VM(vm, "n", int);
  int len = VM(vm, "length", int);
  int maxlen = len;
  if (vm.count("maxlen")) maxlen = VM(vm, "maxlen", int);
  if (maxlen < len) maxlen=len;

  int period = VM(vm, "period", int);
  MSG("++ Simulation output=" << output << " model=" << model
      << " n=" << n << " length=" << len);

  FileAccessRead(in, model);
  vom::VOM m("ACGT", period);
  m.load_node(in);

  FileAccessWrite(out, output);
  for (int i = 0; i < n; i++) {
    out << ">simul " << i  << ENDL;
    std::string s;
    int length = rand() % (maxlen-len+1) + len;
    m.simulate(length, s);
    for (int j = 0; j < sz(s); j++) {
      out << s[j];
      if (j+1 % 80 == 0) out << ENDL;
    }
    if (sz(s) % 80 != 0) out << ENDL;
  }
}

void rand_fom()
{
  std::string output = VM(vm, "output", std::string);
  int order = VM(vm, "order", int);
  int period = VM(vm, "period", int);
  MSG("++ Build Random FOM output=" << output << " order=" << order);
	
  vom::VOM m("ACGT", period);
  m.build_rand_fom(order);
	
  FileAccessWrite(out, output);
  m.print_node(out);
  out.close();
}

void accuracy()
{
  std::string posposFile = VM(vm, "pospos", std::string);
  std::string negposFile = VM(vm, "negpos", std::string);
  std::string posnegFile = VM(vm, "posneg", std::string);
  std::string negnegFile = VM(vm, "negneg", std::string);
  std::string output = VM(vm, "output", std::string);
  MSG("++ Accuracy pospos=" << posposFile << " posneg=" << posnegFile
      << " negneg=" << negnegFile << " negpos=" << negposFile);

  FileAccessRead(pospos, posposFile);
  FileAccessRead(negpos, negposFile);
  FileAccessRead(posneg, posnegFile);
  FileAccessRead(negneg, negnegFile);

  // positive set
  int line = 1;
  double n = 0, n_pos = 0, n_correct = 0, n_pos_correct = 0;

  while (!pospos.eof()) {
    FileAccessLine(pospos, posposAnno);
    FileAccessLine(posneg, posnegAnno);
    //MSG("pp=" << posposAnno << " pn=" << posnegAnno);
    if (posposAnno.empty() || posnegAnno.empty()) {
      MSG("positive set detect end of file line=" << line);
      break;
    }
    if (posposAnno[0] != '>' || posnegAnno[0] != '>') 
      FATAL_ERROR("annotation not found line=" << line);
    line++;

    FileAccessLine(pospos, posposLLStr);
    FileAccessLine(posneg, posnegLLStr);
    
    double posposLL = RAND::s2type<double>(posposLLStr);
    double posnegLL = RAND::s2type<double>(posnegLLStr);

    n++;
    n_pos++;
    if (posposLL > posnegLL) {
      n_correct++;
      n_pos_correct++;
    }
    line++;
  }

  line = 1;
  while (!negneg.eof()) {
    FileAccessLine(negpos, negposAnno);
    FileAccessLine(negneg, negnegAnno);
    if (negposAnno.empty() || negnegAnno.empty()) {
      MSG("negative set detect end of file line=" << line);
      break;
    }
    if (negposAnno[0] != '>' || negnegAnno[0] != '>') 
      FATAL_ERROR("annotation not found line=" << line);
    line++;

    FileAccessLine(negpos, negposLLStr);
    FileAccessLine(negneg, negnegLLStr);
    
    double negposLL = RAND::s2type<double>(negposLLStr);
    double negnegLL = RAND::s2type<double>(negnegLLStr);

    n++;
    if (negnegLL >= negposLL) {
      n_correct++;
    }
    line++;
  }

  FileAccessWrite(out, output);
  MSG(n << "\t" << n_pos << "\t" << n-n_pos);
  out << std::fixed << std::setprecision(2);
  out << (n_correct/n)*100 << "\t" << (n_pos_correct / n_pos)*100 << "\t" << (n_correct-n_pos_correct) / (n-n_pos) * 100 << ENDL;
}

void depth()
{
  std::string output = VM(vm, "output", std::string);
  std::string model = VM(vm, "model", std::string);
  int period = VM(vm, "period", int);

  MSG("++ Depth output=" << output << " model=" << model);

  FileAccessRead(in, model);
  vom::VOM m("ACGT", period);
  m.load_node(in);

  m.depth();
  FileAccessWrite(out, output);
  out << m.count_ << " " << m.depth_ << ENDL;
}

template <typename T>
std::string strcat(const std::string& s, const T& v) 
{
  std::stringstream ss;
  ss << s << v;
  return ss.str();
}

void split()
{
  std::string input = VM(vm, "input", std::string);
  int length = VM(vm, "length", int);
  std::string output = strcat(strcat(input, "."), length);
  if (vm.count("output")) output = VM(vm, "output", std::string);
  MSG("++ Split to short length input=" << input << " output=" << output
      << " length=" << length);

  Sequences seqs;
  seqs.fasta_read(input);

  FileAccessWrite(out, output);
  seqs.split(out, length);
}

void split_train_test() 
{
  std::string input = VM(vm, "input", std::string);
  std::string train = input + ".train";
  std::string test = input + ".test";
  if (vm.count("train")) train = VM(vm, "train", std::string);
  if (vm.count("test")) test = VM(vm, "test", std::string);
  double ratio = VM(vm, "ratio", double);

  MSG("++ Split train test input=" << input << " ratio=" << ratio
      << " train=" << train << " test=" << test);

  Sequences seqs, trainSeqs, testSeqs;
  seqs.fasta_read(input);
  MSG("n_total=" << sz(seqs));

  for (int i = 0; i < sz(seqs); i++) {
    double r = arma::randu();
    if (r < ratio) { // to train set
      trainSeqs.pb(seqs[i]);
      trainSeqs.annotation_.pb(seqs.annotation_[i]);
    }
    else { // to test set
      testSeqs.pb(seqs[i]);
      testSeqs.annotation_.pb(seqs.annotation_[i]);
    }
  }
  MSG("n_train=" << sz(trainSeqs));
  MSG("n_test=" << sz(testSeqs));

  FileAccessWrite(trainout, train);
  FileAccessWrite(testout, test);
  trainSeqs.print(trainout, true);
  testSeqs.print(testout, true);
}

void fix_cds() {
  std::string input = VM(vm, "input", std::string);
  std::string output = input + ".fixed";
  if (vm.count("output")) output = VM(vm, "output", std::string);
  MSG("++ Fix CDS sequences input=" << input << " output=" << output);
	
  Sequences seqs, outSeqs;
  seqs.fasta_read(input);	
  MSG("n_seqs=" << sz(seqs));

  std::vector<vom::sastring> train_set;
  BOOST_FOREACH(const std::string& s, seqs)
    train_set.pb(s);

  int count_codon = 0, count_end_codon = 0, wrong_length = 0, count_wrong_end_codon = 0;
  for (int i = 0; i < sz(seqs); i++) {
    int length = sz(seqs[i]);
    if (length % 3 != 0) {
      MSG(">" << seqs.annotation_[i] << ENDL << " length=" << length << " not divisible by 3");
      wrong_length++;
      continue; // CDS length divisible by 3
    }
    std::string start_codon = seqs[i].substr(0,3);
    if (start_codon != "ATG") {
      MSG(">" << seqs.annotation_[i] << ENDL  << " alternative start codon=" << start_codon);
      count_codon++;
    }
    
    std::string end_codon = seqs[i].substr(length-3,3);
    if (end_codon != "TAA" && end_codon != "TAG" && end_codon != "TGA") {
      MSG(">" << seqs.annotation_[i] << ENDL << " wrong end codon=" << end_codon);
      count_wrong_end_codon++;
      continue;
    }

    bool ok = true;
    int left, size;
    train_set[i].sasearch("TAA", left, size);
    for (int j = left; j < left+size; j++) {
      int pos = train_set[i].SA[j];
      if (pos % 3 == 0 && pos < length - 3) {
	MSG(">" << seqs.annotation_[i] << ENDL << " wrong end TAA at pos=" << pos << " " << seqs[i].substr(pos,3));
	ok = false;
	count_end_codon++;
	break;
      }
    }
    if (!ok) continue;

    train_set[i].sasearch("TAG", left, size);
    for (int j = left; j < left+size; j++) {
      int pos = train_set[i].SA[j];
      if (pos % 3 == 0 && pos < length - 3) {
	MSG(">" << seqs.annotation_[i] << ENDL  << " wrong end TAG at pos=" << pos << " " << seqs[i].substr(pos,3));
	ok = false;
	count_end_codon++;
	break;
      }
    }
    if (!ok) continue;

    train_set[i].sasearch("TGA", left, size);
    for (int j = left; j < left+size; j++) {
      int pos = train_set[i].SA[j];
      if (pos % 3 == 0 && pos < length - 3) {
	MSG(">" << seqs.annotation_[i] << ENDL  << " wrong end TGA at pos=" << pos << " " << seqs[i].substr(pos,3));
	ok = false;
	count_end_codon++;
	break;
      }
    }
    if (!ok) continue;
    
    outSeqs.pb(seqs[i]);
    outSeqs.annotation_.pb(seqs.annotation_[i]);
    
  }
  
  MSG("n_seqs=" << sz(seqs));
  MSG("n_wrong_length=" << wrong_length);
  MSG("n_alternative_start = " << count_codon);
  MSG("n_wrong_end_codon = " << count_wrong_end_codon);
  MSG("n_wrong_inside_end_codon = " << count_end_codon);
  MSG("n_remain_correct = " << sz(outSeqs));
  FileAccessWrite(out, output);
  outSeqs.print(out, true);
}

void gc_content() {
  std::vector<std::string> inputs;
  if (vm.count("inputs")) inputs = VM(vm, "inputs", std::vector<std::string>);
  if (vm.count("input")) inputs.pb(VM(vm, "input", std::string));
  std::string output = VM(vm, "output", std::string);

  int n = 0, n_gc = 0;
  BOOST_FOREACH(const std::string& input, inputs) {
    Sequences seqs;
    seqs.fasta_read(input);
    n += seqs.n;
    n_gc += seqs.n_gc;
  }
  
  FileAccessWrite(out, output);
  out << std::fixed << std::setprecision(2) << (double)n_gc / n * 100 << ENDL;
}
